{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MXNet/Gluon CNN example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import math\n",
    "import mxnet as mx\n",
    "from mxnet import nd, autograd\n",
    "from mxnet import gluon\n",
    "from common.params import *\n",
    "from common.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Force one-gpu\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OS:  linux\n",
      "Python:  3.6.3 |Anaconda custom (64-bit)| (default, Oct 13 2017, 12:02:49) \n",
      "[GCC 7.2.0]\n",
      "MXNet:  1.3.0\n",
      "Numpy:  1.13.3\n",
      "GPU:  ['Tesla V100-SXM2-16GB', 'Tesla V100-SXM2-16GB', 'Tesla V100-SXM2-16GB', 'Tesla V100-SXM2-16GB']\n",
      "CUDA Version 9.1.85\n",
      "CuDNN Version  7.1.3\n"
     ]
    }
   ],
   "source": [
    "print(\"OS: \", sys.platform)\n",
    "print(\"Python: \", sys.version)\n",
    "print(\"MXNet: \", mx.__version__)\n",
    "print(\"Numpy: \", np.__version__)\n",
    "print(\"GPU: \", get_gpu_name())\n",
    "print(get_cuda_version())\n",
    "print(\"CuDNN Version \", get_cudnn_version())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(n_classes=N_CLASSES):\n",
    "    net = gluon.nn.HybridSequential()\n",
    "    with net.name_scope():\n",
    "        net.add(gluon.nn.Conv2D(channels=50, kernel_size=3, padding=1, activation='relu'))\n",
    "        net.add(gluon.nn.Conv2D(channels=50, kernel_size=3, padding=1))\n",
    "        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n",
    "        net.add(gluon.nn.Activation('relu'))\n",
    "        # Equiv to gluon.nn.LeakyReLU(0)\n",
    "        net.add(gluon.nn.Dropout(0.25))\n",
    "        net.add(gluon.nn.Conv2D(channels=100, kernel_size=3, padding=1, activation='relu'))\n",
    "        net.add(gluon.nn.Conv2D(channels=100, kernel_size=3, padding=1))\n",
    "        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n",
    "        net.add(gluon.nn.Activation('relu'))\n",
    "        net.add(gluon.nn.Dropout(0.25))\n",
    "        net.add(gluon.nn.Flatten())\n",
    "        net.add(gluon.nn.Dense(512, activation='relu'))\n",
    "        net.add(gluon.nn.Dropout(0.25))\n",
    "        net.add(gluon.nn.Dense(n_classes))\n",
    "    return net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_model(net, ctx, lr=LR, momentum=MOMENTUM):\n",
    "    net.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)\n",
    "    trainer = gluon.Trainer(\n",
    "        net.collect_params(), \n",
    "        'sgd',\n",
    "        {'learning_rate': lr, 'momentum':momentum})\n",
    "    criterion = gluon.loss.SoftmaxCrossEntropyLoss()\n",
    "    return trainer, criterion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing train set...\n",
      "Preparing test set...\n",
      "(50000, 3, 32, 32) (10000, 3, 32, 32) (50000,) (10000,)\n",
      "float32 float32 int32 int32\n",
      "CPU times: user 776 ms, sys: 568 ms, total: 1.34 s\n",
      "Wall time: 2.34 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Data into format for library\n",
    "x_train, x_test, y_train, y_test = cifar_for_library(channel_first=True)\n",
    "print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)\n",
    "print(x_train.dtype, x_test.dtype, y_train.dtype, y_test.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4 ms, sys: 0 ns, total: 4 ms\n",
      "Wall time: 3.4 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ctx = mx.gpu()\n",
    "net = build_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8 ms, sys: 0 ns, total: 8 ms\n",
      "Wall time: 4.42 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "trainer, criterion = init_model(net, ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch   0: loss: 1.8582\n",
      "Epoch   1: loss: 1.3819\n",
      "Epoch   2: loss: 1.1333\n",
      "Epoch   3: loss: 0.9515\n",
      "Epoch   4: loss: 0.8145\n",
      "Epoch   5: loss: 0.7097\n",
      "Epoch   6: loss: 0.6174\n",
      "Epoch   7: loss: 0.5324\n",
      "Epoch   8: loss: 0.4575\n",
      "Epoch   9: loss: 0.3964\n",
      "CPU times: user 49 s, sys: 13.3 s, total: 1min 2s\n",
      "Wall time: 37.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "net.hybridize()\n",
    "for j in range(EPOCHS):\n",
    "    train_loss = nd.zeros(1, ctx=ctx)\n",
    "    for i, (data, target) in enumerate(yield_mb(x_train, y_train, BATCHSIZE, shuffle=True)):\n",
    "        # Get samples\n",
    "        data = nd.array(data).as_in_context(ctx)\n",
    "        target = nd.array(target).as_in_context(ctx)\n",
    "        with autograd.record():\n",
    "            # Forwards\n",
    "            output = net(data)\n",
    "            # Loss\n",
    "            loss = criterion(output, target)\n",
    "        # Back-prop\n",
    "        loss.backward()\n",
    "        trainer.step(data.shape[0])\n",
    "        train_loss += loss.mean()\n",
    "    # Log    \n",
    "    print('Epoch %3d: loss: %5.4f'%(j, train_loss.asscalar()/(i+1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 296 ms, sys: 48 ms, total: 344 ms\n",
      "Wall time: 278 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Main evaluation loop: 453ms\n",
    "n_samples = (y_test.shape[0]//BATCHSIZE)*BATCHSIZE\n",
    "y_guess = np.zeros(n_samples, dtype=np.int)\n",
    "y_truth = y_test[:n_samples]\n",
    "c = 0\n",
    "for data, target in yield_mb(x_test, y_test, BATCHSIZE):\n",
    "    # Get samples\n",
    "    data = nd.array(data).as_in_context(ctx)\n",
    "    # Forwards\n",
    "    output = net(data)\n",
    "    pred = nd.argmax(output, axis=1)\n",
    "    # Collect results\n",
    "    y_guess[c*BATCHSIZE:(c+1)*BATCHSIZE] = pred.asnumpy()\n",
    "    c += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy:  0.765324519231\n"
     ]
    }
   ],
   "source": [
    "print(\"Accuracy: \", 1.*sum(y_guess == y_truth)/len(y_guess))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
